"""
初始化脚本，用于设置环境和启动服务
"""
import os
import sys
import argparse
import shutil
from dotenv import load_dotenv

def setup_environment():
    """设置环境"""
    print("正在设置环境...")
    
    # 检查.env文件是否存在，如果不存在则从示例创建
    if not os.path.exists(".env"):
        if os.path.exists(".env.example"):
            shutil.copy(".env.example", ".env")
            print("已从.env.example创建.env文件")
        else:
            print("警告: 未找到.env.example文件，请手动创建.env文件")
    else:
        print(".env文件已存在")
    
    # 加载环境变量
    load_dotenv()
    
    # 检查ROCm环境
    try:
        import torch
        if torch.cuda.is_available():
            device_count = torch.cuda.device_count()
            print(f"检测到 {device_count} 个CUDA/ROCm设备:")
            for i in range(device_count):
                print(f"  设备 {i}: {torch.cuda.get_device_name(i)}")
        else:
            print("警告: 未检测到CUDA/ROCm设备，将使用CPU模式")
    except ImportError:
        print("警告: 未安装PyTorch，请先安装PyTorch")
    
    print("环境设置完成")

def install_dependencies():
    """安装依赖"""
    print("正在安装依赖...")
    
    # 检查requirements.txt文件是否存在
    if not os.path.exists("requirements.txt"):
        print("错误: 未找到requirements.txt文件")
        return False
    
    # 安装依赖
    import subprocess
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"])
        print("依赖安装完成")
        return True
    except subprocess.CalledProcessError:
        print("错误: 依赖安装失败")
        return False

def start_server():
    """启动服务器"""
    print("正在启动服务器...")
    
    # 导入并启动服务器
    try:
        from src.api.app import start_server
        start_server()
    except ImportError as e:
        print(f"错误: 无法启动服务器 - {str(e)}")
        return False
    
    return True

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="ROCm大模型推理API初始化脚本")
    parser.add_argument("--setup", action="store_true", help="设置环境")
    parser.add_argument("--install", action="store_true", help="安装依赖")
    parser.add_argument("--start", action="store_true", help="启动服务器")
    parser.add_argument("--all", action="store_true", help="执行所有操作")
    
    args = parser.parse_args()
    
    # 如果没有指定参数，则显示帮助信息
    if not (args.setup or args.install or args.start or args.all):
        parser.print_help()
        return
    
    # 执行操作
    if args.setup or args.all:
        setup_environment()
    
    if args.install or args.all:
        install_dependencies()
    
    if args.start or args.all:
        start_server()

if __name__ == "__main__":
    main()